// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/prestosql/aggregates/geometric_mean_aggregate.h>
#include <cmath>
#include <pollux/exec/simple_aggregate_adapter.h>
#include <pollux/functions/prestosql/aggregates/aggregate_names.h>

using namespace kumo::pollux::exec;

namespace kumo::pollux::aggregate::prestosql {
    namespace {
        template<typename TInput, typename TResult>
        class GeometricMeanAggregate {
        public:
            using InputType = Row<TInput>;

            using IntermediateType =
            Row</*logSum*/ double,
                /*count*/ int64_t>;

            using OutputType = TResult;

            static bool toIntermediate(
                exec::out_type<Row<double, int64_t> > &out,
                exec::arg_type<TInput> in) {
                out.copy_from(std::make_tuple(std::log(in), 1));
                return true;
            }

            struct AccumulatorType {
                // Sum of `log(value)`.
                double logSum_{0};
                // Count of values.
                int64_t count_{0};

                AccumulatorType() = delete;

                explicit AccumulatorType(
                    HashStringAllocator * /*allocator*/,
                    GeometricMeanAggregate<TInput, TResult> * /*fn*/) {
                }

                void addInput(
                    HashStringAllocator * /*allocator*/,
                    exec::arg_type<TInput> data) {
                    logSum_ += std::log(data);
                    count_ = checkedPlus<int64_t>(count_, 1);
                }

                void combine(
                    HashStringAllocator * /*allocator*/,
                    exec::arg_type<Row<double, int64_t> > other) {
                    POLLUX_CHECK(other.at<0>().has_value());
                    POLLUX_CHECK(other.at<1>().has_value());
                    logSum_ += other.at<0>().value();
                    count_ = checkedPlus<int64_t>(count_, other.at<1>().value());
                }

                bool writeFinalResult(exec::out_type<OutputType> &out) {
                    out = std::exp(logSum_ / count_);
                    return true;
                }

                bool writeIntermediateResult(exec::out_type<IntermediateType> &out) {
                    out = std::make_tuple(logSum_, count_);
                    return true;
                }
            };
        };
    } // namespace

    void registerGeometricMeanAggregate(
        const std::string &prefix,
        bool withCompanionFunctions,
        bool overwrite) {
        const std::string name = prefix + kGeometricMean;

        std::vector<std::shared_ptr<exec::AggregateFunctionSignature> > signatures;

        for (const auto &inputType: {"bigint", "double"}) {
            signatures.push_back(exec::AggregateFunctionSignatureBuilder()
                .returnType("double")
                .intermediateType("row(double,bigint)")
                .argumentType(inputType)
                .build());
        }

        // Register for real input type.
        signatures.push_back(exec::AggregateFunctionSignatureBuilder()
            .returnType("real")
            .intermediateType("row(double,bigint)")
            .argumentType("real")
            .build());

        exec::registerAggregateFunction(
            name,
            std::move(signatures),
            [name](
        core::AggregationNode::Step step,
        const std::vector<TypePtr> &argTypes,
        const TypePtr &resultType,
        const core::QueryConfig & /*config*/)
        -> std::unique_ptr<exec::Aggregate> {
                POLLUX_USER_CHECK_EQ(argTypes.size(), 1, "{} takes one argument", name);
                auto inputType = argTypes[0];

                switch (inputType->kind()) {
                    case TypeKind::BIGINT:
                        return std::make_unique<SimpleAggregateAdapter<
                            GeometricMeanAggregate<int64_t, double> > >(
                            step, argTypes, resultType);
                    case TypeKind::DOUBLE:
                        return std::make_unique<
                            SimpleAggregateAdapter<GeometricMeanAggregate<double, double> > >(
                            step, argTypes, resultType);
                    case TypeKind::REAL:
                        return std::make_unique<
                            SimpleAggregateAdapter<GeometricMeanAggregate<float, float> > >(
                            step, argTypes, resultType);
                    default:
                        POLLUX_USER_FAIL(
                            "Unknown input type for {} aggregation {}",
                            name,
                            inputType->toString());
                }
            },
            {false /*orderSensitive*/, false /*companionFunction*/},
            withCompanionFunctions,
            overwrite);
    }
} // namespace kumo::pollux::aggregate::prestosql
